"""synthetic_structsim.py

    Utilities for generating certain graph shapes.
"""
import math
import networkx as nx
import numpy as np
import pdb
# Following GraphWave's representation of structural similarity

# clique(start, nb_nodes, nb_to_remove=0, role_start=0)：
# 创建一个完全图（clique），其节点数量由 nb_nodes 指定。
# start 参数指定了节点的起始索引。
# role_start 参数指定了角色 ID 的起始索引。
# 可选参数 nb_to_remove 指定要移除的边的数量。
def clique(start, nb_nodes, nb_to_remove=0, role_start=0):
    """ Defines a clique (complete graph on nb_nodes nodes,
    with nb_to_remove  edges that will have to be removed),
    index of nodes starting at start
    and role_ids at role_start
    INPUT:
    -------------
    start       :    starting index for the shape
    nb_nodes    :    int correspondingraph to the nb of nodes in the clique
    role_start  :    starting index for the roles
    nb_to_remove:    int-- numb of edges to remove (unif at RDM)
    OUTPUT:
    -------------
    graph       :    a house shape graph, with ids beginning at start
    roles       :    list of the roles of the nodes (indexed starting at
                     role_start)
    """
    # 创建一个形状为 (nb_nodes, nb_nodes) 的全为 1 的矩阵 a，用于表示一个完全图（clique），其中 nb_nodes 表示节点的数量。
    a = np.ones((nb_nodes, nb_nodes))
    # 将矩阵 a 的对角线元素设为 0，因为在完全图中节点与自身不连接。
    np.fill_diagonal(a, 0)
    # 利用 NetworkX 库中的 from_numpy_matrix 函数，根据矩阵 a 创建一个图 graph，其中矩阵 a 中的元素表示边的连接情况。
    graph = nx.from_numpy_array(a)
    # 获取图 graph 的边列表。
    edge_list = graph.edges().keys()
    # 创建一个长度为 nb_nodes 的列表 roles，其中每个元素都为 role_start，表示节点的角色 ID。
    roles = [role_start] * nb_nodes
    # 检查是否需要移除边
    if nb_to_remove > 0:
        # 随机选择要移除的边的索引，len(edge_list) 表示边列表的长度，nb_to_remove 表示要移除的边的数量。
        lst = np.random.choice(len(edge_list), nb_to_remove, replace=False)
        print(edge_list, lst)
        # 根据随机选择的索引，从边列表中获取要移除的边。
        to_delete = [edge_list[e] for e in lst]
        # 从图 graph 中移除要删除的边
        graph.remove_edges_from(to_delete)
        # 遍历要移除的边的索引
        for e in lst:
            print(edge_list[e][0])
            print(len(roles))
            # 增加与被移除边连接的节点的角色 ID。
            roles[edge_list[e][0]] += 1
            roles[edge_list[e][1]] += 1
    # 创建一个字典 mapping_graph，将节点索引映射到从 start 开始的新索引。
    mapping_graph = {k: (k + start) for k in range(nb_nodes)}
    # 重新标记图 graph 的节点，使其从 start 开始。
    graph = nx.relabel_nodes(graph, mapping_graph)
    # 返回生成的图 graph 和角色 ID 列表 roles。
    return graph, roles

# 定义了一个名为 cycle 的函数，它有三个参数：start 表示形状的起始索引，len_cycle 表示环图的长度，role_start 表示角色 ID 的起始索引，默认为 0。
def cycle(start, len_cycle, role_start=0):
    """Builds a cycle graph, with index of nodes starting at start
    and role_ids at role_start
    INPUT:
    -------------
    start       :    starting index for the shape
    role_start  :    starting index for the roles
    OUTPUT:
    -------------
    graph       :    a house shape graph, with ids beginning at start
    roles       :    list of the roles of the nodes (indexed starting at
                     role_start)
    """
    # 创建一个空的无向图对象 graph，用于表示环图。
    graph = nx.Graph()
    # 向图 graph 中添加节点，节点索引从 start 开始，到 start + len_cycle - 1 结束。
    graph.add_nodes_from(range(start, start + len_cycle))
    # 遍历环图中的节点，不包括最后一个节点
    for i in range(len_cycle - 1):
        # 向图 graph 中添加边，连接节点 start + i 和 start + i + 1，构成环图的环。
        graph.add_edges_from([(start + i, start + i + 1)])
    # 向图 graph 中添加最后一个节点到第一个节点的边，构成环图的闭环。
    graph.add_edges_from([(start + len_cycle - 1, start)])
    # 创建一个长度为 len_cycle 的列表 roles，其中每个元素都为 role_start，表示节点的角色 ID。
    roles = [role_start] * len_cycle
    # 返回生成的环图 graph 和角色 ID 列表 roles。
    return graph, roles


# 定义了一个名为 tree 的函数，它有四个参数：start 表示形状的起始索引，height 表示树的高度，r 表示每个节点的分支数量，默认为 10，role_start 表示角色 ID 的起始索引，默认为 0。
def tree(start, height, r=10, role_start=0):
    """Builds a balanced r-tree of height h
    INPUT:
    -------------
    start       :    starting index for the shape
    height      :    int height of the tree 
    r           :    int number of branches per node 
    role_start  :    starting index for the roles
    OUTPUT:
    -------------
    graph       :    a tree shape graph, with ids beginning at start
    roles       :    list of the roles of the nodes (indexed starting at role_start)
    """
    # 利用 NetworkX 库中的 balanced_tree 函数，根据给定的参数 r 和 height 创建一个平衡 r-ary 树，存储在图对象 graph 中。
    graph = nx.balanced_tree(r, height)
    # 创建一个长度为图 graph 中节点数的列表 roles，其中每个元素都为 0，表示节点的角色 ID。
    roles = [0] * graph.number_of_nodes()
    # 返回生成的树形图 graph 和角色 ID 列表 roles。
    return graph, roles


# 定义了一个名为 ba 的函数，它有四个参数：start 表示形状的起始索引，width 表示图的大小，role_start 表示角色 ID 的起始索引，默认为 0，m 表示每个新增节点要连接到现有节点的数量，默认为 5。
def ba(start, width, role_start=0, m=5):
    """Builds a BA preferential attachment graph, with index of nodes starting at start
    and role_ids at role_start
    INPUT:
    -------------
    start       :    starting index for the shape
    width       :    int size of the graph
    role_start  :    starting index for the roles
    OUTPUT:
    -------------
    graph       :    a house shape graph, with ids beginning at start
    roles       :    list of the roles of the nodes (indexed starting at
                     role_start)
    """
    # 利用 NetworkX 库中的 barabasi_albert_graph 函数，根据给定的参数 width 和 m 创建一个 BA 网络，存储在图对象 graph 中。
    graph = nx.barabasi_albert_graph(width, m)
    # 向图 graph 中添加节点，节点索引从 start 开始，到 start + width - 1 结束。
    graph.add_nodes_from(range(start, start + width))
    # 获取图 graph 中的节点，并按节点索引排序，存储在列表 nids 中。
    nids = sorted(graph)
    # 创建一个字典 mapping，将节点索引重新映射到从 start 开始的新索引。
    mapping = {nid: start + i for i, nid in enumerate(nids)}
    # 重新标记图 graph 的节点，使其从 start 开始。
    graph = nx.relabel_nodes(graph, mapping)
    # 创建一个长度为 width 的列表 roles，其中每个元素都为 role_start，表示节点的角色 ID。
    roles = [role_start for i in range(width)]
    # 返回生成的 BA 网络 graph 和角色 ID 列表 roles。
    return graph, roles


# 定义了一个名为 diamond 的函数，它有两个参数：start 表示形状的起始索引，role_start 表示角色 ID 的起始索引，默认为 0。
def diamond(start, role_start=0):
    # 定义了环图的长度为 6。
    len_cycle = 6
    # 创建一个空的无向图对象 graph，用于表示钻石形状的图。
    graph = nx.Graph()
    # 向图 graph 中添加节点，节点索引从 start 开始，到 start + len_cycle - 1 结束。
    graph.add_nodes_from(range(start, start + len_cycle))
    # 遍历环图中的节点，不包括最后一个节点。
    for i in range(len_cycle - 1):
        # 向图 graph 中添加边，连接节点 start + i 和 start + i + 1，构成环图的环。
        graph.add_edges_from([(start + i, start + i + 1)])
    # 向图 graph 中添加最后一个节点到第一个节点的边，构成环图的闭环。
    graph.add_edges_from([(start + len_cycle - 1, start)])
    # 向图 graph 中添加边，连接最后一个节点和第二个节点，形成钻石形状的边。
    graph.add_edges_from([(start + len_cycle - 1, start + 1)])
    # 向图 graph 中添加边，连接倒数第二个节点和第三个节点，形成钻石形状的边。
    graph.add_edges_from([(start + len_cycle - 2, start + 2)])
    # 创建一个长度为 len_cycle 的列表 roles，其中每个元素都为 role_start，表示节点的角色 ID。
    roles = [role_start] * len_cycle
    # 返回生成的钻石形状的图 graph 和角色 ID 列表 roles。
    return graph, roles

# def diamond(start, role_start=0, r=10):
#     """Builds a diamond graph, with index of nodes starting at start
#     and role_ids at role_start
#     INPUT:
#     -------------
#     start       :    starting index for the shape
#     role_start  :    starting index for the roles
#     OUTPUT:
#     -------------
#     graph       :    a house shape graph, with ids beginning at start
#     roles       :    list of the roles of the nodes (indexed starting at
#                      role_start)
#     """
#     graph = nx.Graph()
#     graph.add_nodes_from(range(start, start + 6))
#     graph.add_edges_from(
#         [
#             (start, start + 1),
#             (start + 1, start + 2),
#             (start + 2, start + 3),
#             (start + 3, start),
#         ]
#     )
#     graph.add_edges_from(
#         [
#             (start + 4, start),
#             (start + 4, start + 1),
#             (start + 4, start + 2),
#             (start + 4, start + 3),
#         ]
#     )
#     graph.add_edges_from(
#         [
#             (start + 5, start),
#             (start + 5, start + 1),
#             (start + 5, start + 2),
#             (start + 5, start + 3),
#         ]
#     )
#     roles = [role_start] * 6
#     return graph, roles

# 定义了一个名为 house 的函数，它有两个参数：start 表示形状的起始索引，role_start 表示角色 ID 的起始索引，默认为 0。
def house(start, role_start=0):
    """Builds a house-like  graph, with index of nodes starting at start
    and role_ids at role_start
    INPUT:
    -------------
    start       :    starting index for the shape
    role_start  :    starting index for the roles
    OUTPUT:
    -------------
    graph       :    a house shape graph, with ids beginning at start
    roles       :    list of the roles of the nodes (indexed starting at
                     role_start)
    """
    # 创建一个空的无向图对象 graph，用于表示房子形状的图。
    graph = nx.Graph()
    # 向图 graph 中添加 5 个节点，节点索引从 start 开始，到 start + 4 结束。
    graph.add_nodes_from(range(start, start + 5))
    # 向图 graph 中添加边，连接房子形状的四条边，构成一个正方形的外框。
    graph.add_edges_from(
        [
            (start, start + 1),
            (start + 1, start + 2),
            (start + 2, start + 3),
            (start + 3, start),
        ]
    )
    # graph.add_edges_from([(start, start + 2), (start + 1, start + 3)])
    # 向图 graph 中添加两条额外的边，连接房子形状的顶点和底部两个顶点，构成房顶和底部。
    graph.add_edges_from([(start + 4, start), (start + 4, start + 1)])
    # 创建一个长度为 5 的列表 roles，其中前两个元素为 role_start，表示房子的外框，接下来两个元素为 role_start + 1，表示房子的顶点，最后一个元素为 role_start + 2，表示房子的底部。
    roles = [role_start, role_start, role_start + 1, role_start + 1, role_start + 2]
    # 返回生成的房子形状的图 graph 和角色 ID 列表 roles。
    return graph, roles


# 定义了一个名为 grid 的函数，它有四个参数：start 表示形状的起始索引，dim 表示网格的维度，r 表示每个节点的分支数量，默认为 10，role_start 表示角色 ID 的起始索引，默认为 0。
def grid(start, dim, r=10, role_start=0):
    """ Builds a 2by2 grid
    """
    # 使用 NetworkX 库中的 grid_graph 函数创建一个二维网格图 grid_G，其中参数 [3, 2] 表示网格的维度。
    grid_G = nx.grid_graph([3, 2])
    # 将网格图 grid_G 的节点标签转换为整数，并从 start 开始标记节点，保持节点之间的连接关系不变。
    grid_G = nx.convert_node_labels_to_integers(grid_G, first_label=start)
    # 创建一个长度为网格图 grid_G 中节点数的列表 roles，其中每个元素都为 role_start，表示节点的角色 ID。
    roles = [role_start for i in grid_G.nodes()]
    # 返回生成的网格图 grid_G 和角色 ID 列表 roles。
    return grid_G, roles

# 定义了一个名为 build_graph 的函数，它有多个参数，用于创建具有特定形状的图，并在基础形状上随机附加形状。
# build_graph 函数的参数，分别表示基础形状的宽度、基础形状的类型、要附加的形状列表、起始节点的索引、是否随机放置形状、要添加的随机边数、每个新节点要连接到现有节点的数量。
def build_graph(
    width_basis,
    basis_type,
    list_shapes,
    start=0,
    rdm_basis_plugins=False,
    add_random_edges=0,
    m=5,
):
    """This function creates a basis (scale-free, path, or cycle)
    and attaches elements of the type in the list randomly along the basis.
    Possibility to add random edges afterwards.
    INPUT:
    --------------------------------------------------------------------------------------
    width_basis      :      width (in terms of number of nodes) of the basis # 300
    basis_type       :      (torus, string, or cycle)                        # ba
    shapes           :      list of shape list (1st arg: type of shape,      # house
                            next args:args for building the shape,
                            except for the start)
    start            :      initial nb for the first node                       # 0
    rdm_basis_plugins:      boolean. Should the shapes be randomly placed       # False
                            along the basis (True) or regularly (False)?
    add_random_edges :      nb of edges to randomly add on the structure        # 0
    m                :      number of edges to attach to existing node (for BA graph) # 5
    OUTPUT:
    --------------------------------------------------------------------------------------
    basis            :      a nx graph with the particular shape
    role_ids         :      labels for each role
    plugins          :      node ids with the attached shapes
    """
    # 如果基础形状类型是 "ba"（即 BA 网络），则执行以下代码块。
    if basis_type == "ba":
        # 调用与基础形状类型对应的函数，例如 ba() 函数，根据给定的参数创建基础形状，并将基础形状图存储在 basis 变量中，将节点角色标签存储在 role_id 变量中。
        basis, role_id = eval(basis_type)(start, width_basis, m=m)
    else:
        # 调用与基础形状类型对应的函数，根据给定的参数创建基础形状，并将基础形状图存储在 basis 变量中，将节点角色标签存储在 role_id 变量中。
        basis, role_id = eval(basis_type)(start, width_basis, r=m)
    # 计算基础形状图中的节点数和要附加的形状数量。
    n_basis, n_shapes = nx.number_of_nodes(basis), len(list_shapes)
    # 更新起始节点索引，以便从基础形状的末尾开始为附加形状分配节点索引。
    start += n_basis  # indicator of the id of the next node

    # Sample (with replacement) where to attach the new motifs
    # 如果 rdm_basis_plugins 参数为 True，则随机选择基础形状图中的节点作为附加形状的起始节点
    if rdm_basis_plugins is True:
        # 从基础形状图的节点中随机选择与要附加的形状数量相同数量的节点索引，这些节点将用作附加形状的起始节点。
        plugins = np.random.choice(n_basis, n_shapes, replace=False)
    else:
        # 否则，将附加形状等间隔地放置在基础形状上。
        # 计算基础形状中节点的间隔，以便将附加形状等间隔地放置在基础形状上。n_basis 是基础形状图中的节点数量，n_shapes 是要附加的形状数量。
        spacing = math.floor(n_basis / n_shapes) # n_shapes: 80
        # 使用列表推导式创建一个列表 plugins，其中包含 n_shapes 个元素，每个元素都是附加形状的起始节点索引。索引值等间隔地分布在基础形状图中，起始索引从0开始，每个索引之间的间隔为 spacing。 int() 函数用于确保计算结果是整数。
        plugins = [int(k * spacing) for k in range(n_shapes)]
    # 创建一个字典 seen_shapes，用于存储每种形状的索引范围。初始时，将基础形状的索引范围存储在其中。
    seen_shapes = {"basis": [0, n_basis]} # n_basis=300
    # 遍历要附加的形状列表，获取形状的索引和形状类型。
    for shape_id, shape in enumerate(list_shapes):
        # 获取当前形状的类型。
        shape_type = shape[0]
        # 创建一个包含起始节点索引的参数列表。
        args = [start]
        # 如果形状列表的长度大于 1，说明还有其他参数。
        if len(shape) > 1:
            # 将附加形状的其他参数添加到参数列表中。
            args += shape[1:]
        # 将初始角色标签索引（默认为 0）添加到参数列表中。
        args += [0]
        # 根据当前形状的类型和参数调用相应的形状函数，并将形状图和节点角色标签存储在 graph_s 和 roles_graph_s 变量中。
        graph_s, roles_graph_s = eval(shape_type)(*args)
        # 计算附加形状图中的节点数。
        n_s = nx.number_of_nodes(graph_s)
        try:
            # 获取当前形状类型在 seen_shapes 字典中的索引范围起始值。
            col_start = seen_shapes[shape_type][0]
        except:
            # 将当前形状的角色标签索引设置为当前角色标签列表中的最大值加一。这确保了新形状的角色标签索引不会与已有的角色标签索引重叠。
            col_start = np.max(role_id) + 1
            # 将当前形状类型的索引范围 [col_start, n_s] 存储到 seen_shapes 字典中，以便在下一次迭代时使用。
            seen_shapes[shape_type] = [col_start, n_s]
        # Attach the shape to the basis
        # 向基础形状图中添加附加形状图的节点。
        basis.add_nodes_from(graph_s.nodes())
        # 向基础形状图中添加附加形状图的边。
        basis.add_edges_from(graph_s.edges())
        # 将附加形状的起始节点与基础形状图中的相应节点连接起来。
        basis.add_edges_from([(start, plugins[shape_id])]) # attach
        # if shape_type == "cycle":
        #     if np.random.random() > 0.5:
        #         a = np.random.randint(1, 4)
        #         b = np.random.randint(1, 4)
        #         basis.add_edges_from([(a + start, b + plugins[shape_id])])
        # 计算附加形状节点的角色标签，将其与当前形状的角色标签索引偏移，以确保角色标签不会与其他形状的角色标签重叠。
        temp_labels = [r + col_start for r in roles_graph_s]
        # temp_labels[0] += 100 * seen_shapes[shape_type][0]
        # 将计算得到的附加形状节点的角色标签添加到总体角色标签列表中。
        role_id += temp_labels
        # 更新起始节点索引，以便下一个附加形状从上一个形状的末尾开始。
        start += n_s
    # 如果要添加的随机边数大于 0
    if add_random_edges > 0:
        # add random edges between nodes:
        # 循环添加随机边的次数
        for p in range(add_random_edges):
            # 随机选择基础形状图中的两个节点作为随机边的起始节点和目标节点。
            src, dest = np.random.choice(nx.number_of_nodes(basis), 2, replace=False)
            print(src, dest)
            # 向基础形状图中添加随机边，将起始节点和目标节点连接起来。
            basis.add_edges_from([(src, dest)])
    # 返回生成的具有特定形状的图 basis、角色标签列表 role_id 和附加形状节点列表 plugins。
    return basis, role_id, plugins
